# tools 基于faiss检索
from langchain.agents import tool, Tool, create_react_agent, AgentExecutor,a
from langchain.chains.retrieval_qa.base import RetrievalQA
from langchain.embeddings import CacheBackedEmbeddings
from langchain.storage import LocalFileStore
# 导入缓存聊天
from langchain_community.chat_message_histories import ChatMessageHistory
from langchain_community.document_loaders import TextLoader
from langchain_community.embeddings import DashScopeEmbeddings
from langchain_community.llms.tongyi import Tongyi
from langchain_community.vectorstores import FAISS
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables.history import RunnableWithMessageHistory
from langchain_text_splitters import CharacterTextSplitter

llm = Tongyi()
embeddings = DashScopeEmbeddings()
store = LocalFileStore("./cache/")

cache_embedder = cache_embeddings = CacheBackedEmbeddings.from_bytes_store(
    embeddings,
    store,
    namespace=embeddings.model,
)
doc = TextLoader('new.txt', encoding='utf-8').load()

spliter = CharacterTextSplitter("\n", chunk_size=200, chunk_overlap=0)
chunks = spliter.split_documents(doc)
db = FAISS.from_documents(chunks, cache_embedder)
retriever = db.as_retriever()  # 解释：检索器
faq_chain = RetrievalQA.from_chain_type(
    llm=llm,
    chain_type="stuff",
    retriever=retriever,
)


@tool('FAQ')
def faq(user_input):
    """
    当你需要回答关于商品问题的时候,比如咨询等问题,使用当前函数的返回结果
    """
    return faq_chain.invoke(user_input)


def search_order(user_input):
    print('订单号', user_input)
    if user_input == '000000':
        return '订单号' + '000000'
    return '订单,失效'


def recommend_goods(user_input):
    print('商品推荐', user_input)
    return '商品推荐' + '华为手机，小米手机，oppo手机'


gtool = [
    Tool(
        name="search_order",
        description="用于查询订单号",
        func=search_order
    ),
    Tool(
        name="recommend_goods",
        description="用于推荐商品",
        func=recommend_goods
    ),
    faq

]

template = """
Answer the following questions as best you can. You have access to the following tools,请不要添加自己的回答,只返回工具返回内容即可,不要添加自己的文字:

            {tools}

            Use the following format:

            Question: the input question you must answer
            Thought: you should always think about what to do
            Action: the action to take, should be one of [{tool_names}]
            Action Input: the input to the action
            Observation: the result of the action
            ... (this Thought/Action/Action Input/Observation can repeat N times)
            Thought: I now know the final answer
            Final Answer: the final answer to the original input question

            Begin!

            Question: {input}
            Thought:{agent_scratchpad}
"""
prompt_template = ChatPromptTemplate.from_template(template)

agent = create_react_agent(
    tools=gtool,
    llm=llm,
    prompt=prompt_template,
)
agent_executor = AgentExecutor(
    agent=agent,
    tools=gtool,
    verbose=True,
)
message_history = ChatMessageHistory()

agent_with_chat_history = RunnableWithMessageHistory(
    agent_executor,
    lambda session_id: message_history,
    input_messages_key="input",
    history_messages_key="chat_history",
)
res = agent_with_chat_history.invoke(
    {"input": "邮寄到三亚需要多久"},
    config={
        "configurable": {
            "session_id": "session-10086",
        }
    }
)